use polars::frame::column::Column;
use anyhow::{Context, Result};
use calamine::{Reader, Xlsx, open_workbook, Data as CalDataType};
use polars::prelude::*;
use std::path::Path;

// https://github.com/ToucanToco/fastexcel/tree/main/src/utils
#[derive(Debug)]
pub struct ExcelReader {
    converters: Vec<(String, DataType)>,
}

impl ExcelReader {
    pub fn new(converters: &[(&str, DataType)]) -> Self {
        Self {
            converters: converters
                .iter()
                .map(|(name, dtype)| (name.to_string(), dtype.clone()))
                .collect(),
        }
    }

    pub fn read_excel(&self, path: &Path, sheet_index: usize) -> Result<DataFrame> {
        // 打开 Excel 文件
        let mut workbook: Xlsx<_> = open_workbook(path)?;

        // 选择工作表
        let range = workbook
            .worksheet_range_at(sheet_index)
            .ok_or_else(|| anyhow::anyhow!("Sheet index {} not found", sheet_index))??;

        // 处理表头
        let headers: Vec<String> = range
            .rows()
            .next()
            .ok_or_else(|| anyhow::anyhow!("No header row"))?
            .iter()
            .map(|cell| cell.to_string())
            .collect();

        // 初始化列数据存储
        let mut columns: Vec<Vec<Option<String>>> = vec![vec![]; headers.len()];

        // 处理数据行
        for row in range.rows().skip(1) {
            for (col_idx, cell) in row.iter().enumerate() {
                let value = self.process_cell(cell, &headers[col_idx]);
                columns[col_idx].push(value);
            }
        }

        // 转换为 Polars DataFrame
        let columns: Vec<Column> = headers
            .iter()
            .enumerate()
            .map(|(idx, name)| {
                let dtype = self
                    .converters
                    .iter()
                    .find(|(cname, _)| cname == name)
                    .map(|(_, dtype)| dtype)
                    .unwrap_or(&DataType::String);

                let data = &columns[idx];
                let series = match dtype {
                    DataType::String => Series::new(
                        name.as_str().into(),
                        data.iter()
                            .map(|v| v.as_deref().unwrap_or(""))
                            .collect::<Vec<&str>>(),
                    ),
                    DataType::Int64 => Series::new(
                        name.as_str().into(),
                        data
                            .iter()
                            .map(|v| v.as_deref().and_then(|s| s.parse().ok()))
                            .collect::<Vec<Option<i64>>>(),
                    ),
                    DataType::Float64 => Series::new(
                        name.as_str().into(),
                        data
                            .iter()
                            .map(|v| v.as_deref().and_then(|s| s.parse().ok()))
                            .collect::<Vec<Option<f64>>>(),
                    ),
                    _ => panic!("Unsupported data type: {:?}", dtype),
                };
                Column::new(
                    name.as_str().into(),
                    series
                )
            })
            .collect();

        DataFrame::new(columns).context("Failed to create DataFrame")
    }

    fn process_cell(&self, cell: &CalDataType, column_name: &str) -> Option<String> {
        use CalDataType::*; // 添加这行引入枚举变体
        let converter = self.converters.iter().find(|(name, _)| name == column_name);

        match (cell, converter) {
            // 处理需要特殊转换的列
            (Float(f), Some((_, DataType::String))) => {
                Some(format!("{:.0}", f)) // 身份证号等数字转字符串
            }
            (Int(i), Some((_, DataType::String))) => {
                Some(i.to_string()) // 整数转字符串
            }
            // 通用处理
            (Empty, _) => None,
            (String(s), _) => Some(s.clone()),
            (Float(f), _) => Some(f.to_string()),
            (Int(i), _) => Some(i.to_string()),
            (Bool(b), _) => Some(b.to_string()),
            (DateTime(_), _) => Some(cell.to_string()),
            _ => Some("".to_string()),
        }
    }
}
